import json
import re
from rag_chain import RAGChain
from text_processor import TextVectorizer

# 目标电话号码和公司
target_phone = "15937098754"
expected_company = "北京科技创新有限公司"

# 测试数据
test_data = """
这是一个测试文件，包含电话号码与公司的关联信息。

以下是一些公司和联系人信息：
1. 北京科技创新有限公司 - 张三 - 15937098754
2. 上海未来信息技术有限公司 - 李四 - 13812345678
3. 广州智慧城市发展有限公司 - 王五 - 13987654321
"""

# 初始化测试
print("=== 开始最终电话号码检索测试 ===")

# 1. 初始化向量器和RAG链
vectorizer = TextVectorizer(use_tfidf=True)
rag_chain = RAGChain(deepseek_api_key="sk-a6d2da9cbeaf4cf59512a993de770cbd", vectorizer=vectorizer)
print("已初始化向量器和RAG链")

# 2. 手动添加测试数据到向量器
vectorizer.vectorize_and_store(test_data, metadata={"source": "test_phone_company.txt"})
print("已手动添加测试数据到向量器")

# 3. 执行RAG链查询
print("\n=== 执行RAG链查询 ===")
results = rag_chain.run(target_phone, n_results=3, similarity_metric="cosine")
print(f"RAG链结果: {json.dumps(results, ensure_ascii=False, indent=2)}")

# 4. 验证结果
print("\n=== 结果验证 ===")
if "answer" in results:
    print(f"AI回答: {results['answer']}")
    if expected_company in results['answer'] and target_phone in results['answer']:
        print(f"成功! AI回答正确识别了{target_phone}属于{expected_company}")
    else:
        # 检查引用文档中是否包含关联信息
        found_in_references = False
        if "references" in results:
            for ref in results['references']:
                if expected_company in ref.get('full_text', '') and target_phone in ref.get('full_text', ''):
                    found_in_references = True
                    print(f"在引用文档中找到关联信息: {expected_company} - {target_phone}")
                    print(f"但AI回答未正确提取，可能需要进一步优化提示模板。")
                    break
        if not found_in_references:
            print(f"未找到{target_phone}与{expected_company}的关联信息")
else:
    print("RAG链结果格式不正确")

print("\n=== 测试完成 ===")